"""
python fujian1_autoArima_findArgs.py
"""

import os
import json
from pmdarima import auto_arima
import pandas as pd

# Input and output paths
input_dir = "fujian/fujian1/cubic_spline/interpolation_output"
output_file_path = "fujian/fujian1/auto_arima/args/all_args.json"
os.makedirs(os.path.dirname(output_file_path), exist_ok=True)

# Initialize a list to store all ARIMA parameters
all_args = []

# Iterate through each JSON file
for filename in os.listdir(input_dir):
    if filename.endswith(".json"):
        filepath = os.path.join(input_dir, filename)
        
        # Read JSON data
        with open(filepath, 'r', encoding='utf-8') as file:
            data = json.load(file)

        # Convert to DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)

        # Use auto_arima to automatically select ARIMA parameters
        auto_model = auto_arima(df['inventory'],
                                start_p=0, max_p=2,
                                start_q=0, max_q=2,
                                d=None,              # Automatically choose differencing order
                                seasonal=False,      # Non-seasonal data
                                stepwise=True,       # Enable efficient stepwise algorithm
                                suppress_warnings=True)
        
        # Get the best (p, d, q) parameters
        best_p, best_d, best_q = auto_model.order

        # Extract category ID from filename
        category_id = filename.split("category")[-1].split(".")[0]

        # Append the parameters to the list
        all_args.append({
            "category_id": category_id,
            "p": best_p,
            "d": best_d,
            "q": best_q
        })

        print(f"Processed category {category_id}: p={best_p}, d={best_d}, q={best_q}")

# Write all parameters to a single JSON file
with open(output_file_path, 'w', encoding='utf-8') as output_file:
    json.dump(all_args, output_file, ensure_ascii=False, indent=4)

print(f"All parameters saved to: {output_file_path}")
